package com.dahuyou.addtrycatch;

import com.dahuyou.addtrycatch.filter.MethodRecordingFilter;
import com.sun.corba.se.spi.logging.CORBALogDomains;
import org.objectweb.asm.*;
import org.objectweb.asm.commons.AdviceAdapter;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;

import static org.objectweb.asm.Opcodes.ASM5;

public class AddTryCatchTransformer implements ClassFileTransformer {
    /**
     * 所有的类被来加载器加载的时候都会走这个方法，这样我们就有机会通过各种技术手段来对字节码进行插桩了
     *
     * @param loader
     * @param className           类全限定名称
     * @param classBeingRedefined
     * @param protectionDomain
     * @param classfileBuffer     类字节码对应的二进制数组，可以通过defineClass直接加载到JVM中并生成Class对象
     * @return
     */
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        // 过滤不需要插桩的类
        if (MethodRecordingFilter.isNotNeedInject(className)) {
            return classfileBuffer;
        }
        byte[] bytes = new byte[0];
        try {
            bytes = getBytes(className);
        } catch (IOException e) {
            e.printStackTrace();
        }
        // 写到磁盘的.class文件中，方便调试
        outputClazz(bytes, className);
        return bytes;
    }

    private byte[] getBytes(String className) throws IOException {
        ClassReader cr = new ClassReader(className);
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);
        cr.accept(new ClassVisitor(ASM5, cw) {
            public MethodVisitor visitMethod(int access, String name, String
                    descriptor, String signature, String[] exceptions) {
                // 方法过滤
                if (!"strToNumber".equals(name))
                    return super.visitMethod(access, name, descriptor,
                            signature, exceptions);
                MethodVisitor mv = super.visitMethod(access, name, descriptor,
                        signature, exceptions);
                return new AdviceAdapter(ASM5, mv, access, name, descriptor) {
                    private Label from = new Label(); // try 开始位置
                    private Label to = new Label(); // try 结束位置
                    private Label target = new Label(); // catch开始位置
                    // 方法进入时修改字节码
                    protected void onMethodEnter() {
                        //标志：try块开始位置 以下是ASM提供好的模板代码，这样写就行，定义了位置信息以及要catch的异常信息
                        visitLabel(from);
                        visitTryCatchBlock(from,
                                to,
                                target,
                                "java/lang/Exception");
                    }

                    // 访问局部变量和操作数栈 ???
                    // 貌似是访问在物理上结束时调用
                    public void visitMaxs(int maxStack, int maxLocals) {
                        //标志：try块结束
                        mv.visitLabel(to);

                        // catch块开始
                        mv.visitLabel(target);
                        mv.visitFrame(Opcodes.F_SAME1, 0, null, 1, new Object[] {"java/lang/Exception"});

                        // 创建catch(Exception e) 中的e并设置其值为当前栈顶的异常对象，此时局部变量表就有了e异常对象了
                        // long是位置信息吧
                        int local = newLocal(Type.LONG_TYPE);
                        mv.visitVarInsn(ASTORE, local); // 准备异常对象到local位置的局部变量中，这样后续要使用的话就直接拿了

                        mv.visitLdcInsn(className + "." + name);
                        mv.visitVarInsn(ALOAD, local);
                        mv.visitMethodInsn(INVOKESTATIC, Type.getInternalName(LogUtil.class), "point", "(Ljava/lang/String;Ljava/lang/Throwable;)V", false);

                        // 加载局部变量表中的异常对象到操作数栈
                        mv.visitVarInsn(ALOAD, local);
                        // throw栈顶的异常
                        mv.visitInsn(ATHROW);
                        super.visitMaxs(maxStack, maxLocals);
                    }

                    // 方法退出时修改字节码
                    protected void onMethodExit(int opcode) {
                        if ((IRETURN <= opcode && opcode <= RETURN) || opcode == ATHROW) {

//                        System.out.println("this.nextLocal: " + this.nextLocal);
                            // 注意：此时栈顶的值就是要return回去的那个值
                            // astore	将栈顶引用类型数值存入指定本地变量
                            // aload	将指定的引用类型本地变量推送至栈顶

                            int nextLocal = this.nextLocal; // 获取可用的最近的局部变量表的位置，这一步就相当于定义了变量 Integer var
                            mv.visitVarInsn(ASTORE, nextLocal); // 将当前栈顶元素的值赋值到新创建的局部变量 相当于Integer var = 栈顶的整数值
                            mv.visitVarInsn(ALOAD, nextLocal); // 加载局部变量到栈顶
//                        System.out.println("this.nextLocal: " + this.nextLocal); // next local值就变了
//                            mv.visitInsn(ARETURN); return貌似是自动加的，这里显示加的话，就会报一个areturn指令的错误

                            // 准备com.dahuyou.addtrycatch.LogUtil.point方法需要的参数，即都压到栈顶去
                            mv.visitLdcInsn(className + "." + name);
                            mv.visitVarInsn(ALOAD, nextLocal);
                            // 调用静态方法com.dahuyou.addtrycatch.LogUtil.point
                            mv.visitMethodInsn(INVOKESTATIC, Type.getInternalName(LogUtil.class), "point", "(Ljava/lang/String;Ljava/lang/Object;)V", false);
                        }
                    }
                };
            }
        }, ClassReader.EXPAND_FRAMES);
        return cw.toByteArray();
    }

    private void outputClazz(byte[] bytes, String className) {
        // 输出类字节码
        FileOutputStream out = null;
        try {
            String pathName = AddTryCatchTransformer.class.getResource("/").getPath() + className + "_after_instrument.class";
            out = new FileOutputStream(new File(pathName));
            System.out.println("插桩后代码输出路径：" + pathName);
            out.write(bytes);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (null != out) try {
                out.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

}
